/*
This is part of OpenLoong Dynamics Control, an open project for the control of biped robot,
Copyright (C) 2024-2025 Humanoid Robot (Shanghai) Co., Ltd.
Feel free to use in any purpose, and cite OpenLoong-Dynamics-Control in any style, to contribute to the advancement of the community.
 <https://atomgit.com/openloong/openloong-dyn-control.git>
 <web@openloong.org.cn>
*/
#include "MJ_interface.h"
#include "GLFW_callbacks.h"
#include <string>
#include <iostream>
#include "data_logger.h"
#include "PVT_ctrl.h"
#include "pino_kin_dyn.h"
#include "useful_math.h"
#include "data_bus.h"

const double dt = 0.001;
// MuJoCo load and compile model
char error[1000] = "Could not load binary model";
mjModel *mj_model = mj_loadXML("../models/scene.xml", 0, error, 1000);
mjData *mj_data = mj_makeData(mj_model);

int main(int argc, char **argv) {
    // initialize classes
    UIctr uiController(mj_model, mj_data);   // UI control for Mujoco
    MJ_Interface mj_interface(mj_model, mj_data); // data interface for Mujoco
    Pin_KinDyn kinDynSolver("../models/AzureLoong.urdf"); // kinematics and dynamics solver
    PVT_Ctr pvtCtr(mj_model->opt.timestep, "../common/joint_ctrl_config.json");// PVT joint control
    DataBus RobotState(kinDynSolver.model_nv); // data bus
    DataLogger logger("../record/stand_datalog.log"); // data logger
    int model_nv = kinDynSolver.model_nv;

    // initialize UI: GLFW
    uiController.iniGLFW();
    uiController.enableTracking(); // enable viewpoint tracking of the body 1 of the robot
    uiController.createWindow("Stand Demo", true);

    // ini data logger
    logger.addIterm("simTime", 1);
    logger.addIterm("motor_pos_des", model_nv - 6);
    logger.addIterm("motor_pos_cur", model_nv - 6);
    logger.addIterm("motor_vel_cur", model_nv - 6);
    logger.addIterm("motor_tor_des", model_nv - 6);
    logger.addIterm("motor_tor_out", model_nv - 6);
    logger.addIterm("rpyVal", 3);
    logger.addIterm("gpsVal", 3);
    logger.addIterm("fe_l_pos_L_des", 3);
    logger.addIterm("fe_r_pos_L_des", 3);
    logger.addIterm("fe_l_pos_W", 3);
    logger.addIterm("fe_r_pos_W", 3);
    logger.finishItermAdding();

    // ini position and posture for foot-end and hand
    Eigen::Vector3d fe_l_pos_L_des = {-0.018, 0.113, -1.01};
    Eigen::Vector3d fe_r_pos_L_des = {-0.018, -0.116, -1.01};
    Eigen::Vector3d fe_l_eul_L_des = {-0.000, -0.01, -0.000};
    Eigen::Vector3d fe_r_eul_L_des = {0.000, -0.01, 0.000};
    Eigen::Matrix3d fe_l_rot_des = eul2Rot(fe_l_eul_L_des(0), fe_l_eul_L_des(1), fe_l_eul_L_des(2));
    Eigen::Matrix3d fe_r_rot_des = eul2Rot(fe_r_eul_L_des(0), fe_r_eul_L_des(1), fe_r_eul_L_des(2));

    Eigen::Vector3d hd_l_pos_L_des = {-0.02, 0.32, -0.159};
    Eigen::Vector3d hd_r_pos_L_des = {-0.02, -0.32, -0.159};
    Eigen::Vector3d hd_l_eul_L_des = {-1.253, 0.122, -1.732};
    Eigen::Vector3d hd_r_eul_L_des = {1.253, 0.122, 1.732};
    Eigen::Matrix3d hd_l_rot_des = eul2Rot(hd_l_eul_L_des(0), hd_l_eul_L_des(1), hd_l_eul_L_des(2));
    Eigen::Matrix3d hd_r_rot_des = eul2Rot(hd_r_eul_L_des(0), hd_r_eul_L_des(1), hd_r_eul_L_des(2));
    
    // 计算初始站立姿态
    auto resLeg = kinDynSolver.computeInK_Leg(fe_l_rot_des, fe_l_pos_L_des, fe_r_rot_des, fe_r_pos_L_des);
    auto resHand = kinDynSolver.computeInK_Hand(hd_l_rot_des, hd_l_pos_L_des, hd_r_rot_des, hd_r_pos_L_des);

    // 设置仿真参数
    double simTime = mj_data->time;
    double simEndTime = 5;
    double gravityCompensation = 1.0; // 重力补偿系数

    // Main loop:
    while (!glfwWindowShouldClose(uiController.window)) {
        mjtNum simstart = mj_data->time;
        while (mj_data->time - simstart < 1.0 / 60.0 && uiController.runSim) {
            mj_step(mj_model, mj_data);
            simTime = mj_data->time;
            
            // Read the sensors:
            mj_interface.updateSensorValues();
            mj_interface.dataBusWrite(RobotState);

            // update kinematics and dynamics info
            kinDynSolver.dataBusRead(RobotState);
            kinDynSolver.computeJ_dJ();
            kinDynSolver.computeDyn();
            kinDynSolver.dataBusWrite(RobotState);

            // 站立控制逻辑 - 保持初始姿势
            fe_l_pos_L_des = RobotState.fe_l_pos_L;
            fe_r_pos_L_des = RobotState.fe_r_pos_L;
            RobotState.motors_pos_des = eigen2std(resLeg.jointPosRes + resHand.jointPosRes);
            RobotState.motors_vel_des.assign(model_nv - 6, 0);
            
            // 添加重力补偿
            RobotState.motors_tor_des.assign(model_nv - 6, 0);
            for (int i = 0; i < model_nv - 6; i++) {
                RobotState.motors_tor_des[i] = gravityCompensation * RobotState.dyn_Non[i + 6];
            }
            
            // 设置腿部状态为支撑相
            RobotState.legState = DataBus::DSt;
            
            // joint PVT controller - 启用位置控制
            pvtCtr.dataBusRead(RobotState);
            for (int i = 0; i < model_nv - 6; i++) {
                pvtCtr.enablePV(i);
            }
            pvtCtr.calMotorsPVT();
            pvtCtr.dataBusWrite(RobotState);

            // 发送关节力矩命令
            mj_interface.setMotorsTorque(RobotState.motors_tor_out);
            
            // 记录关键数据
            logger.startNewLine();
            logger.recItermData("simTime", simTime);
            logger.recItermData("motor_pos_des", RobotState.motors_pos_des);
            logger.recItermData("motor_pos_cur", RobotState.motors_pos_cur);
            logger.recItermData("motor_vel_cur", RobotState.motors_vel_cur);
            logger.recItermData("motor_tor_des", RobotState.motors_tor_des);
            logger.recItermData("motor_tor_out", RobotState.motors_tor_out);
            logger.recItermData("rpyVal", RobotState.rpy);
            logger.recItermData("gpsVal", RobotState.base_pos);
            logger.recItermData("fe_l_pos_L_des", fe_l_pos_L_des);
            logger.recItermData("fe_r_pos_L_des", fe_r_pos_L_des);
            logger.recItermData("fe_l_pos_W", RobotState.fe_l_pos_W);
            logger.recItermData("fe_r_pos_W", RobotState.fe_r_pos_W);
            logger.finishLine();
        };
        
        if (mj_data->time >= simEndTime) {
            std::cout << "5秒站立完成，例程结束" << std::endl;
            break;
        }

        uiController.updateScene();
    };

    // 关闭窗口
    uiController.Close();

    return 0;
}
